import torch
from torch_geometric.nn import GATConv

class GAT_Model(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, num_heads=1):
        super().__init__()
        # 使用 GATConv 替代 DevConvGNN
        self.gat1 = GATConv(in_channels, hidden_channels, heads=num_heads)
        self.gat2 = GATConv(hidden_channels * num_heads, hidden_channels, heads=num_heads)
        self.gat3 = GATConv(hidden_channels * num_heads, out_channels, heads=1, concat=False)

    def forward(self, graph):
        edge_index, x = graph.edge_index.long(), graph.x

        x = self.gat1(x, edge_index)  # Apply GATConv layer 1
        x = torch.relu(x)  # Activation after the first layer

        x = self.gat2(x, edge_index)  # Apply GATConv layer 2
        x = torch.relu(x)  # Activation after the second layer

        x = self.gat3(x, edge_index)  # Apply GATConv layer 3
        x = torch.sigmoid(x)  # Activation after the last layer

        return x
